﻿using System;
using System.Collections.Generic;
using System.Data.Services.Client;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net;
using System.Text;
using System.Threading;
using System.Xml;
using System.Xml.Linq;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;
using Microsoft.WindowsAzure.StorageClient.Protocol;

namespace AntsCode.AzureStorageBackup
{
    public class Restore
    {
        private CloudStorageAccount account;
        private CloudStorageAccount restoreFromAccount;
        private string backupName;
        private int maxThreads;
        private IEnumerable<string> priorityTables;
        private IEnumerable<string> priorityContainers;

        private static XNamespace AtomNamespace = "http://www.w3.org/2005/Atom";
        private static XNamespace AstoriaDataNamespace = "http://schemas.microsoft.com/ado/2007/08/dataservices";
        private static XNamespace AstoriaMetadataNamespace = "http://schemas.microsoft.com/ado/2007/08/dataservices/metadata";

        private readonly string backupPrefix = "backup";

        public Restore(CloudStorageAccount account, 
            CloudStorageAccount restoreFromAccount, 
            string restoreFromBackupName,
             int maxThreads = 5,
            IEnumerable<string> priorityTables = null,
            IEnumerable<string> priorityContainers = null)
        {
            this.account = account;
            this.restoreFromAccount = restoreFromAccount;
            this.backupName = string.Format("{0}-{1}", backupPrefix, restoreFromBackupName);
            this.maxThreads = maxThreads;
            this.priorityTables = priorityTables != null ? priorityTables : new List<string>();
            this.priorityContainers = priorityContainers != null ? priorityContainers : new List<string>();
        }

        public void BeginRestore()
        {
            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Storage restore started.", DateTime.UtcNow);

            // Load the backup data
            var restoreFromBlobClient = StorageUtil.GetBlobClient(this.restoreFromAccount);
            var restoreFromContainer = restoreFromBlobClient.GetContainerReference(this.backupName);
            var backupDataBlob = restoreFromContainer.GetBlobReference("BackupData.xml");

            MemoryStream stream = new MemoryStream();
            backupDataBlob.DownloadToStream(stream);
            stream.Position = 0;

            XmlDocument doc = new XmlDocument();
            doc.Load(stream);

            // Get tables
            var tableNodes = doc.GetElementsByTagName("Table");
            var tables = new List<string>(tableNodes.Count);

            foreach (XmlNode tableNode in tableNodes)
            {
                tables.Add(tableNode.InnerText);
            }

            // Get containers
            var containerNodes = doc.GetElementsByTagName("Container");
            var containers = new List<string>(containerNodes.Count);

            foreach (XmlNode containerNode in containerNodes)
            {
                containers.Add(containerNode.InnerText);
            }

            // Build consolidated restore list
            var restoreList = StorageUtil.CreateRunList(tables, containers, this.priorityTables, this.priorityContainers, RestoreTable, RestoreContainer);
        
            var threadRunner = new ThreadRunner();
            threadRunner.Run(restoreList, this.maxThreads);

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Storage restore completed.", DateTime.UtcNow);
        }

        private void RestoreTable(string tableName)
        {
            var tableClient = StorageUtil.GetTableClient(this.account);
            var restoreFromBlobClient = StorageUtil.GetBlobClient(this.restoreFromAccount);
            var restoreFromContainer = restoreFromBlobClient.GetContainerReference(this.backupName);

            // Delete any existing table
            tableClient.DeleteTableIfExist(tableName);
            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Starting to restore table {1} in thread {2}", DateTime.UtcNow, tableName, Thread.CurrentThread.ManagedThreadId);

            // Re-create the table
            CreateTableWithRetries(tableClient, tableName);

            BlobRequestOptions options = new BlobRequestOptions()
            {
                RetryPolicy = RetryPolicies.RetryExponential(5, RetryPolicies.DefaultClientBackoff)
            };

            // Look for backed up data for this table
            var directory = restoreFromContainer.GetDirectoryReference(string.Format("TableData/{0}", tableName));

            var blobRefs = directory.ListBlobs();
            CloudBlockBlob blob;

            foreach (var blobRef in blobRefs)
            {
                blob = restoreFromContainer.GetBlockBlobReference(blobRef.Uri.AbsoluteUri);

                // get all blocks and for each block read it separately as it is an xml doc by itself
                var blocks = (IEnumerable<ListBlockItem>)blob.DownloadBlockList(BlockListingFilter.Committed, options);

                long currentOffset = 0;
                foreach (ListBlockItem block in blocks)
                {
                    // read each block using the range
                    HttpWebRequest request = BlobRequest.Get(blob.Uri, 120, null, currentOffset, block.Size, null);
                    blob.ServiceClient.Credentials.SignRequest(request);

                    XDocument doc = null;

                    using (HttpWebResponse response = (HttpWebResponse)request.GetResponse())
                    {
                        using (Stream stream = response.GetResponseStream())
                        {
                            doc = XDocument.Load(new XmlTextReader(stream));
                        }
                    }

                    IEnumerable<XElement> batchNodes = doc.Element("Block").Elements("Batch");
                    foreach (XElement batchNode in batchNodes)
                    {
                        ExecuteBatch(tableClient, tableName, blob, batchNode);
                    }

                    currentOffset += block.Size;
                }
            }

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Restored table {1}.", DateTime.UtcNow, tableName);
        }

        private void CreateTableWithRetries(CloudTableClient tableClient, string tableName)
        {
            bool success = false;

            while (!success)
            {
                try
                {
                    tableClient.CreateTable(tableName);
                    success = true;
                }
                catch (StorageClientException)
                {
                    // The table has not finished deleting, so wait
                    Thread.Sleep(5000);
                }
            }
        }

        static void ExecuteBatch(CloudTableClient tableClient, string tableName, CloudBlockBlob blob, XElement batchNode)
        {
            TableServiceContext context = tableClient.GetDataServiceContext();
            context.WritingEntity += new EventHandler<ReadingWritingEntityEventArgs>(OnWritingEntity);

            // for each entry create a backup entity
            IEnumerable<XElement> entries = batchNode.Elements(AtomNamespace + "entry");
            foreach (XElement entryNode in entries)
            {
                XElement propertiesElem = entryNode.Elements(AtomNamespace + "content")
                                                .Elements(AstoriaMetadataNamespace + "properties")
                                                .FirstOrDefault();
                XElement pkElement = propertiesElem.Element(AstoriaDataNamespace + "PartitionKey");
                XElement rkElement = propertiesElem.Element(AstoriaDataNamespace + "RowKey");

                BackupEntity entity = new BackupEntity()
                {
                    PartitionKey = pkElement.Value,
                    RowKey = rkElement.Value,
                    EntryElement = entryNode
                };

                context.AddObject(tableName, entity);
            }

            try
            {
                context.SaveChangesWithRetries(SaveChangesOptions.Batch);
            }
            catch (Exception e)
            {
                var sb = new StringBuilder();
                BackupEntity backupEntity;

                foreach (var entity in context.Entities)
                {
                    if (sb.Length > 0)
                    {
                        sb.Append(", ");
                    }

                    backupEntity = (BackupEntity)entity.Entity;
                    sb.Append(string.Format("{0}/{1}", backupEntity.PartitionKey, backupEntity.RowKey));
                }

                // Log the error & continue
                Trace.TraceWarning(string.Format("Warning: Batch failed with error {0}{1}\n\nBackup blob: {2}\n\nTable: {3}\n\nEntities in batch: {4}", 
                    e.Message, 
                    e.InnerException != null ? "\n" + e.InnerException.Message : "",
                    blob.Uri.AbsoluteUri,
                    tableName,
                    sb.ToString()));
            }
        }

        private void RestoreContainer(string containerName)
        {
            var blobClient = StorageUtil.GetBlobClient(this.account);
            var restoreFromBlobClient = StorageUtil.GetBlobClient(this.restoreFromAccount);
            var restoreFromContainer = restoreFromBlobClient.GetContainerReference(this.backupName);

            if (containerName.StartsWith(this.backupPrefix, StringComparison.CurrentCultureIgnoreCase))
            {
                // Do not restore any backup containers
                Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Container {1} excluded as it is a backup container.", DateTime.UtcNow, containerName);
                return;
            }

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Starting to restore container {1} in thread {2}", DateTime.UtcNow, containerName, Thread.CurrentThread.ManagedThreadId);

            // Delete any existing container
            var container = blobClient.GetContainerReference(containerName);

            try
            {
                container.Delete();
            }
            catch (StorageClientException)
            {
                // There is no "delete if exists"
            }

            // Re-create the container
            CreateContainerWithRetries(blobClient, containerName);

            // Look for backed up data for this container
            var directory = restoreFromContainer.GetDirectoryReference(string.Format("BlobData/{0}", containerName));

            var blobRefs = directory.ListBlobs();

            CloudBlob sourceBlob;
            CloudBlob destBlob;
            string blobName;

            foreach (var blobRef in blobRefs)
            {
                sourceBlob = restoreFromContainer.GetBlockBlobReference(blobRef.Uri.AbsoluteUri);

                // Get the full blob name 
                blobName = sourceBlob.Uri.AbsolutePath.Substring(sourceBlob.Uri.AbsolutePath.IndexOf(containerName) + containerName.Length + 1);

                destBlob = container.GetBlobReference(blobName);
     
                StorageUtil.CopyBlob(sourceBlob, destBlob);
            }

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Restored container {1}", DateTime.UtcNow, containerName);
        }

        private void CreateContainerWithRetries(CloudBlobClient blobClient, string containerName)
        {
            bool success = false;
            var container = blobClient.GetContainerReference(containerName);

            while (!success)
            {
                try
                {
                    container.Create();
                    success = true;
                }
                catch (StorageClientException)
                {
                    // The table has not finished deleting, so wait
                    Thread.Sleep(5000);
                }
            }
        }

        static void OnWritingEntity(object sender, ReadingWritingEntityEventArgs args)
        {
            BackupEntity entity = args.Entity as BackupEntity;
            XElement content = args.Data.Element(AtomNamespace + "content");
            XElement propertiesElem = content.Element(AstoriaMetadataNamespace + "properties");

            propertiesElem.Remove();

            XElement propertiesElemToUse = entity.EntryElement.Elements(AtomNamespace + "content")
                                            .Elements(AstoriaMetadataNamespace + "properties")
                                            .FirstOrDefault();

            content.Add(propertiesElemToUse);
        }
    }
}